import psycopg2
from utils.read_yaml import ReadYaml


class PostgresDB:
    '''
    连接postgresql数据库
    '''

    def __init__(self, host=None, port=None, user=None, password=None, dbname=None):
        self.read_yaml = ReadYaml()
        self.conn = None  # 改名为 conn 而不是 connect
        self.dbname = dbname
        self.user = user
        self.password = password
        self.host = host
        self.port = port
        if not host:
            self.host = self.read_yaml.get_value('dbConfig/host')
        if not port:
            self.port = self.read_yaml.get_value('dbConfig/port')
        if not user:
            self.user = self.read_yaml.get_value('dbConfig/user')
        if not password:
            self.password = self.read_yaml.get_value('dbConfig/password')
        if not dbname:
            self.dbname = self.read_yaml.get_value('dbConfig/database')

    def connect(self):
        """建立到PostgreSQL数据库的连接"""
        try:
            self.conn = psycopg2.connect(  # 使用 conn 而不是 connect
                dbname=self.dbname,
                user=self.user,
                password=self.password,
                host=self.host,
                port=self.port
            )
            return self.conn
        except Exception as e:
            print(f"无法连接到数据库: {e}")
            raise

    def select_data(self, sql):
        if not self.conn:
            self.connect()
        cursor = self.conn.cursor()  # 使用 self.conn 而不是 self.connect
        cursor.execute(sql)
        res = cursor.fetchall()
        cursor.close()
        return res

    def change_data(self, sql):
        if not self.conn:
            self.connect()
        cursor = self.conn.cursor()  # 使用 self.conn 而不是 self.connect
        cursor.execute(sql)
        num = cursor.rowcount
        self.conn.commit()
        cursor.close()
        return f"{sql} \n【操作影响到 {num}条数据】"

    def __call__(self, sql: str, connect=True):
        sql_new = sql.lower().strip()
        start = sql_new.split(" ")[0]
        if start in ["select", 'use', 'show']:
            res = self.select_data(sql)
            if not connect:
                self.conn.close()  # 使用 self.conn 而不是 self.connect
            return res
        elif start in ["delete", "alter", "update", "insert", "drop"]:
            res = self.change_data(sql)
            if not connect:
                self.conn.close()  # 使用 self.conn 而不是 self.connect
            return res


if __name__ == '__main__':
    db = PostgresDB()
    rows = db("SELECT * from user_info where nickname = 'Ljiang5';", connect=False)
    print(rows)